# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

from abc import abstractmethod
from collections import namedtuple
import enum

from sortedcontainers import SortedDict


class SizeUnit(enum.Enum):
  QTY = 1
  AMT = 2


PriceQty = namedtuple('PriceQty', ['price', 'qty'])
NotionalSize = namedtuple('NotionalSize', ['accumulated', 'entries', 'unit'])


class PriceMap(object):
  def __init__(self):
    self._dict = SortedDict()
    self.clear_cache()

  @abstractmethod
  def price_to_key(self, price: float) -> float:
    pass

  @abstractmethod
  def key_to_price(self, key: float) -> float:
    pass

  def clear_cache(self):
    self._cache_min = None
    self._cache_notional_qty = (None, None)
    self._cache_notional_amt = (None, None)

  def clear(self):
    self._dict.clear()

  def get(self, price: float, default=None):
    return self._dict.get(self.price_to_key(price), default)

  def not_empty(self):
    return bool(self._dict)

  def __contains__(self, price: float):
    return self.price_to_key(price) in self._dict

  def __getitem__(self, price: float):
    return self._dict[self.price_to_key(price)]

  def __setitem__(self, price: float, entry):
    self.clear_cache()
    self._dict[self.price_to_key(price)] = entry

  def __delitem__(self, price: float):
    self.clear_cache()
    del self._dict[self.price_to_key(price)]

  def iteritems(self):
    for key, value in self._dict.items():
      yield (self.key_to_price(key), value)

  def get_min(self):
    if self._cache_min is not None:
      return self._cache_min

    if self._dict:
      key, value = self._dict.peekitem(0)
      self._cache_min = value
      return value
    else:
      return None

  def iter_entries(self, from_price, to_price, exclude_max=False):
    from_price = self.price_to_key(from_price)
    to_price = self.price_to_key(to_price)
    min_price = min(from_price, to_price)
    max_price = max(from_price, to_price)
    for key, value in self._dict.irange(min_price, max_price, inclusive=(True, not exclude_max)):
      yield value

  def get_entries(self, from_price, to_price, exclude_max=False):
    return [v for v in self.iter_entries(from_price, to_price, exclude_max)]

  def get_notional_entries_by_qty(self, qty):
    if self._cache_notional_qty[0] == qty:
      return self._cache_notional_qty[1]

    accumulated_qty = 0
    result = []
    for key, value in self.iteritems():
      assert key == value.price
      cur_price = key
      cur_qty = value.qty
      assert cur_qty > 0
      remaining_qty = qty - accumulated_qty
      if cur_qty >= remaining_qty:
        accumulated_qty += remaining_qty
        result.append(PriceQty(cur_price, remaining_qty))
        break
      else:
        accumulated_qty += cur_qty
        result.append(PriceQty(cur_price, cur_qty))

    ret = NotionalSize(accumulated_qty, result, unit=SizeUnit.QTY)
    self._cache_notional_qty = (qty, ret)
    return ret

  def get_notional_entries_by_amt(self, amt):
    if self._cache_notional_amt[0] == amt:
      return self._cache_notional_amt[1]

    accumulated_amt = 0
    result = []
    for key, value in self.iteritems():
      assert key == value.price
      cur_price = key
      cur_qty = value.qty
      assert cur_qty > 0
      remaining_amt = amt - accumulated_amt
      if cur_qty * cur_price >= remaining_amt:
        accumulated_amt += remaining_amt
        result.append(PriceQty(cur_price, remaining_amt / cur_price))
        break
      else:
        accumulated_amt += cur_price * cur_qty
        result.append(PriceQty(cur_price, cur_qty))

    ret = NotionalSize(accumulated_amt, result, unit=SizeUnit.AMT)
    self._cache_notional_amt = (amt, ret)
    return ret


class BidPriceMap(PriceMap):
  def price_to_key(self, price):
    return -price

  def key_to_price(self, key):
    return -key


class AskPriceMap(PriceMap):
  def price_to_key(self, price):
    return price

  def key_to_price(self, key):
    return key
